#include <bsl_err_internal.h>
#include <bsl_errno.h>
#include <bsl_params.h>
#include <crypt_algid.h>
#include <crypt_errno.h>
#include <securec.h>

#include <wd.h>
#include <wd_bmm.h>
#include <wd_digest.h>

#include "kaep_err.h"
#include "kaev1_digests.h"
#include "kaev1_memory.h"

#include <stdint.h>
#include <stdio.h>

KaeQueuePoolHead *gDigestsNodePool = NULL;

int WdDigestsInitQnodePool(void)
{
    KaeQueuePoolDestroy(gDigestsNodePool);

    gDigestsNodePool = KaeInitQueuePool(WCRYPTO_DIGEST);
    if (gDigestsNodePool == NULL) {
        return KAE_WD_DO_DIGEST_FAIL;
    }
    return CRYPT_SUCCESS;
}

int WdDigestsDestroyQnodePool(void)
{
    KaeQueuePoolDestroy(gDigestsNodePool);
    return CRYPT_SUCCESS;
}

void SecureDigestsCallback(const void *msg, void *tag)
{
    if (!msg || !tag) {
        return;
    }
    struct wcrypto_digest_msg *message = (struct wcrypto_digest_msg *)msg;
    KAE_DIGEST_Ctx *ctx = (KAE_DIGEST_Ctx *)tag;
    memcpy(ctx->opData.out, message->out, message->out_bytes);
}

void KAE_DIGEST_Free(KAE_DIGEST_Ctx *ctx)
{
    if (ctx == NULL) {
        return;
    }

    if (ctx->digestCtx) {
        wcrypto_del_digest_ctx(ctx->digestCtx);
        ctx->digestCtx = NULL;
    }

    if (ctx->pool && ctx->opData.in) {
        ctx->setup.br.free(ctx->setup.br.usr, ctx->opData.in);
        ctx->opData.in = NULL;
    }

    if (ctx->pool && ctx->opData.out) {
        ctx->setup.br.free(ctx->setup.br.usr, ctx->opData.out);
        ctx->opData.out = NULL;
    }

    KaePutNodeToPool(gDigestsNodePool, ctx->entry);

    free(ctx);
}

KAE_DIGEST_Ctx *KAE_DIGEST_NewCtx(CRYPT_MD_AlgId md_alg)
{
    KAE_DIGEST_Ctx *ctx = (KAE_DIGEST_Ctx *)malloc(sizeof(KAE_DIGEST_Ctx));
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    (void)memset_s(ctx, sizeof(KAE_DIGEST_Ctx), 0, sizeof(KAE_DIGEST_Ctx));

    ctx->entry = KaeGetNodeFromPool(gDigestsNodePool);
    if (ctx->entry == NULL) {
        KAE_DIGEST_Free(ctx);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    ctx->queue = ctx->entry->queue;
    ctx->pool = ctx->entry->mempool;

    switch (md_alg) {
        case CRYPT_MD_SM3: {
            ctx->digestSize = SM3_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SM3;
            break;
        }
        case CRYPT_MD_MD5: {
            ctx->digestSize = MD5_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_MD5;
            break;
        }
        case CRYPT_MD_SHA1: {
            ctx->digestSize = SHA1_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA1;
            break;
        }
        case CRYPT_MD_SHA256: {
            ctx->digestSize = SHA256_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA256;
            break;
        }
        case CRYPT_MD_SHA224: {
            ctx->digestSize = SHA224_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA224;
            break;
        }
        case CRYPT_MD_SHA384: {
            ctx->digestSize = SHA384_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA384;
            break;
        }
        case CRYPT_MD_SHA512: {
            ctx->digestSize = SHA512_DIGEST_SIZE;
            ctx->setup.alg = WCRYPTO_SHA512;
            break;
        }
        default: {
            BSL_ERR_PUSH_ERROR(KAE_WD_NOT_SUPPORT);
            goto ERR;
        }
    }

    ctx->id = md_alg;
    ctx->setup.mode = WCRYPTO_DIGEST_NORMAL;
    ctx->setup.br.alloc = (void *)wd_alloc_blk;
    ctx->setup.br.free = (void *)wd_free_blk;
    ctx->setup.br.iova_map = (void *)wd_blk_iova_map;
    ctx->setup.br.iova_unmap = (void *)wd_blk_iova_unmap;
    ctx->setup.br.usr = ctx->pool;
    ctx->setup.cb = SecureDigestsCallback;
    ctx->digestCtx = wcrypto_create_digest_ctx(ctx->queue, &ctx->setup);
    if (ctx->digestCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        goto ERR;
    }

    ctx->opData.in = wd_alloc_blk(ctx->pool);
    if (ctx->opData.in == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    ctx->opData.out = wd_alloc_blk(ctx->pool);
    if (ctx->opData.out == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    ctx->opData.out_bytes = ctx->digestSize;
    (void)memset_s(ctx->opData.out, ctx->digestSize, 0, ctx->digestSize);
    ctx->opData.has_next = true;

    return ctx;
ERR:
    KAE_DIGEST_Free(ctx);
    return NULL;
}

int32_t KAE_DIGEST_Init(KAE_DIGEST_Ctx *ctx, BSL_Param *param)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    (void)param;

    return CRYPT_SUCCESS;
}

int32_t WD_Do_DigestImpl(KAE_DIGEST_Ctx *ctx)
{
    int32_t ret;
    int tryCount = 0;

AGAIN:
    ret = wcrypto_do_digest(ctx->digestCtx, &ctx->opData, NULL);
    if (ret != WD_SUCCESS) {
        if (ret == -WD_EBUSY && tryCount <= 5) {
            tryCount++;
            goto AGAIN;
        } else {
            return KAE_WD_DO_DIGEST_FAIL;
        }
    }

    return CRYPT_SUCCESS;
}

int32_t KAE_DIGEST_Update(KAE_DIGEST_Ctx *ctx, const uint8_t *in, uint32_t len)
{
    if ((ctx == NULL) || (ctx->digestCtx == NULL) || (in == NULL && len != 0)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (len == 0) {
        return CRYPT_SUCCESS;
    }

    int cnt = 0;

    if (ctx->remain > 0) {
        if (len + ctx->remain > DIGEST_CACHE_SIZE) {
            (void)memcpy_s(ctx->block + ctx->remain, DIGEST_CACHE_SIZE - ctx->remain, in, DIGEST_CACHE_SIZE - ctx->remain);
            in += DIGEST_CACHE_SIZE - ctx->remain;
            len -= (DIGEST_CACHE_SIZE - ctx->remain);
            ctx->remain = 0;
            ctx->opData.in_bytes = DIGEST_CACHE_SIZE;
            ctx->opData.has_next = true;
            (void)memcpy_s(ctx->opData.in, DIGEST_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
            if (WD_Do_DigestImpl(ctx) != CRYPT_SUCCESS) {
                BSL_ERR_PUSH_ERROR(KAE_WD_DO_DIGEST_FAIL);
                return KAE_WD_DO_DIGEST_FAIL;
            }
            ctx->updatedCount++;
        } else {
            (void)memcpy_s(ctx->block + ctx->remain, DIGEST_CACHE_SIZE - ctx->remain, in, len);
            ctx->remain += len;
            return CRYPT_SUCCESS;
        }
    }

    while (len - cnt > DIGEST_CACHE_SIZE) {
        ctx->opData.in_bytes = DIGEST_CACHE_SIZE;
        ctx->opData.has_next = true;
        (void)memcpy_s(ctx->opData.in, DIGEST_CACHE_SIZE, in, ctx->opData.in_bytes);
        if (WD_Do_DigestImpl(ctx) != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(KAE_WD_DO_DIGEST_FAIL);
            return KAE_WD_DO_DIGEST_FAIL;
        }
        ctx->updatedCount++;
        in += DIGEST_CACHE_SIZE;
        cnt += DIGEST_CACHE_SIZE;
    }

    ctx->remain = len - cnt;
    (void)memcpy_s(ctx->block, DIGEST_CACHE_SIZE, in, ctx->remain);

    return CRYPT_SUCCESS;
}

uint8_t sm3Digest[] = {
    0x1a, 0xb2, 0x1d, 0x83, 0x55, 0xcf, 0xa1, 0x7f, 0x8e, 0x61, 0x19, 0x48, 0x31, 0xe8, 0x1a, 0x8f,
    0x22, 0xbe, 0xc8, 0xc7, 0x28, 0xfe, 0xfb, 0x74, 0x7e, 0xd0, 0x35, 0xeb, 0x50, 0x82, 0xaa, 0x2b,
};

uint8_t md5Digest[] = {
    0xd4, 0x1d, 0x8c, 0xd9, 0x8f, 0x00, 0xb2, 0x04, 0xe9, 0x80, 0x09, 0x98, 0xec, 0xf8, 0x42, 0x7e,
};

uint8_t sha1Digest[] = {
    0xda, 0x39, 0xa3, 0xee, 0x5e, 0x6b, 0x4b, 0x0d, 0x32, 0x55,
    0xbf, 0xef, 0x95, 0x60, 0x18, 0x90, 0xaf, 0xd8, 0x07, 0x09,
};

uint8_t sha224Digest[] = {0xd1, 0x4a, 0x02, 0x8c, 0x2a, 0x3a, 0x2b, 0xc9, 0x47, 0x61, 0x02, 0xbb, 0x28, 0x82,
                          0x34, 0xc4, 0x15, 0xa2, 0xb0, 0x1f, 0x82, 0x8e, 0xa6, 0x2a, 0xc5, 0xb3, 0xe4, 0x2f};

uint8_t sha256Digest[] = {0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14, 0x9a, 0xfb, 0xf4,
                          0xc8, 0x99, 0x6f, 0xb9, 0x24, 0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b,
                          0x93, 0x4c, 0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55};

uint8_t sha384Digest[] = {0x38, 0xb0, 0x60, 0xa7, 0x51, 0xac, 0x96, 0x38, 0x4c, 0xd9, 0x32, 0x7e,
                          0xb1, 0xb1, 0xe3, 0x6a, 0x21, 0xfd, 0xb7, 0x11, 0x14, 0xbe, 0x07, 0x43,
                          0x4c, 0x0c, 0xc7, 0xbf, 0x63, 0xf6, 0xe1, 0xda, 0x27, 0x4e, 0xde, 0xbf,
                          0xe7, 0x6f, 0x65, 0xfb, 0xd5, 0x1a, 0xd2, 0xf1, 0x48, 0x98, 0xb9, 0x5b};

uint8_t sha512Digest[] = {0xcf, 0x83, 0xe1, 0x35, 0x7e, 0xef, 0xb8, 0xbd, 0xf1, 0x54, 0x28, 0x50, 0xd6,
                          0x6d, 0x80, 0x07, 0xd6, 0x20, 0xe4, 0x05, 0x0b, 0x57, 0x15, 0xdc, 0x83, 0xf4,
                          0xa9, 0x21, 0xd3, 0x6c, 0xe9, 0xce, 0x47, 0xd0, 0xd1, 0x3c, 0x5d, 0x85, 0xf2,
                          0xb0, 0xff, 0x83, 0x18, 0xd2, 0x87, 0x7e, 0xec, 0x2f, 0x63, 0xb9, 0x31, 0xbd,
                          0x47, 0x41, 0x7a, 0x81, 0xa5, 0x38, 0x32, 0x7a, 0xf9, 0x27, 0xda, 0x3e};

void KaeWDZeroInputFinal(KAE_DIGEST_Ctx *ctx, uint8_t *out, uint32_t *len)
{
    switch (ctx->id) {
        case CRYPT_MD_SM3: {
            (void)memcpy_s(out, *len, sm3Digest, SM3_DIGEST_SIZE);
            *len = SM3_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_MD5: {
            (void)memcpy_s(out, *len, md5Digest, MD5_DIGEST_SIZE);
            *len = MD5_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_SHA1: {
            (void)memcpy_s(out, *len, sha1Digest, SHA1_DIGEST_SIZE);
            *len = SHA1_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_SHA256: {
            (void)memcpy_s(out, *len, sha256Digest, SHA256_DIGEST_SIZE);
            *len = SHA256_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_SHA224: {
            (void)memcpy_s(out, *len, sha224Digest, SHA224_DIGEST_SIZE);
            *len = SHA224_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_SHA384: {
            (void)memcpy_s(out, *len, sha384Digest, SHA384_DIGEST_SIZE);
            *len = SHA384_DIGEST_SIZE;
            break;
        }
        case CRYPT_MD_SHA512: {
            (void)memcpy_s(out, *len, sha512Digest, SHA512_DIGEST_SIZE);
            *len = SHA512_DIGEST_SIZE;
            break;
        }
        default: {
            BSL_ERR_PUSH_ERROR(KAE_WD_NOT_SUPPORT);
            return;
        }
    }
}

int32_t KAE_DIGEST_Final(KAE_DIGEST_Ctx *ctx, uint8_t *out, uint32_t *len)
{
    if (ctx == NULL || ctx->digestCtx == NULL || (out == NULL) || (len == NULL)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (*len < ctx->digestSize) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_BUFF_LEN_NOT_ENOUGH);
        return CRYPT_EAL_BUFF_LEN_NOT_ENOUGH;
    }

    ctx->opData.in_bytes = ctx->remain;
    ctx->opData.has_next = false;

    if (ctx->remain == 0 && ctx->updatedCount == 0) {
        KaeWDZeroInputFinal(ctx, out, len);
        return CRYPT_SUCCESS;
    }

    (void)memcpy_s(ctx->opData.in, DIGEST_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
    if (WD_Do_DigestImpl(ctx) != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_DO_DIGEST_FAIL);
        return KAE_WD_DO_DIGEST_FAIL;
    }
    (void)memcpy_s(out, ctx->digestSize, ctx->opData.out, ctx->digestSize);
    *len = ctx->digestSize;

    return CRYPT_SUCCESS;
}

int32_t KAE_DIGEST_Deinit(KAE_DIGEST_Ctx *ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    wd_free_blk(ctx->pool, ctx->opData.in);
    wd_free_blk(ctx->pool, ctx->opData.out);

    ctx->opData.in = wd_alloc_blk(ctx->pool);
    if (ctx->opData.in == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }

    ctx->opData.out = wd_alloc_blk(ctx->pool);
    if (ctx->opData.out == NULL) {
        wd_free_blk(ctx->pool, ctx->opData.in);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }

    ctx->opData.out_bytes = ctx->digestSize;
    (void)memset_s(ctx->opData.out, ctx->digestSize, 0, ctx->digestSize);
    ctx->opData.has_next = false;

    (void)memset_s(ctx->block, DIGEST_CACHE_SIZE, 0, DIGEST_CACHE_SIZE);
    ctx->remain = 0;

    return CRYPT_SUCCESS;
}

KAE_DIGEST_Ctx *KAE_DIGEST_Dup(const KAE_DIGEST_Ctx *src)
{
    if (src == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return NULL;
    }

    KAE_DIGEST_Ctx *newCtx = KAE_DIGEST_NewCtx(src->id);
    if (newCtx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    return newCtx;
}